"""
快速启动脚本
提供交互式菜单，方便用户快速使用项目功能
"""
import os
import sys


def print_menu():
    """打印主菜单"""
    print("\n" + "=" * 80)
    print("雷达俯视图点目标实例分割项目 - 快速启动菜单")
    print("=" * 80)
    print("\n请选择操作：")
    print("\n【训练相关】")
    print("  1. 训练单个模型（指定模型和fold）")
    print("  2. 训练单个模型（所有fold）")
    print("  3. 训练所有模型（所有fold）")
    print("\n【推理相关】")
    print("  4. 单个模型推理")
    print("  5. 所有模型推理（对比）")
    print("\n【其他】")
    print("  6. 查看项目结构")
    print("  7. 查看使用说明")
    print("  0. 退出")
    print("\n" + "=" * 80)


def train_single_model_single_fold():
    """训练单个模型的单个fold"""
    print("\n可用模型：unet, deeplabv3plus, fcn, pspnet, unetpp, segnet")
    model_name = input("请输入模型名称: ").strip().lower()

    fold = input("请输入fold编号（0-7）: ").strip()

    epochs = input("训练轮数（默认120，直接回车使用默认值）: ").strip()
    batch_size = input("批大小（默认8，直接回车使用默认值）: ").strip()

    cmd = f"python baselines/{model_name}/train.py --fold {fold}"

    if epochs:
        cmd += f" --epochs {epochs}"
    if batch_size:
        cmd += f" --batch_size {batch_size}"

    print(f"\n执行命令: {cmd}\n")
    os.system(cmd)


def train_single_model_all_folds():
    """训练单个模型的所有fold"""
    print("\n可用模型：unet, deeplabv3plus, fcn, pspnet, unetpp, segnet")
    model_name = input("请输入模型名称: ").strip().lower()

    epochs = input("训练轮数（默认120，直接回车使用默认值）: ").strip()
    batch_size = input("批大小（默认8，直接回车使用默认值）: ").strip()

    cmd = f"python baselines/{model_name}/train.py"

    if epochs:
        cmd += f" --epochs {epochs}"
    if batch_size:
        cmd += f" --batch_size {batch_size}"

    print(f"\n执行命令: {cmd}\n")
    os.system(cmd)


def train_all_models():
    """训练所有模型"""
    print("\n将训练所有6个模型的8折交叉验证")
    confirm = input("这将需要较长时间，确认继续？(y/n): ").strip().lower()

    if confirm == 'y':
        cmd = "python train_all.py"
        print(f"\n执行命令: {cmd}\n")
        os.system(cmd)
    else:
        print("已取消")


def infer_single_model():
    """单个模型推理"""
    print("\n可用模型：unet, deeplabv3plus, fcn, pspnet, unetpp, segnet")
    model_name = input("请输入模型名称: ").strip().lower()

    image_path = input("请输入图像路径: ").strip()
    fold = input("使用哪个fold的模型（默认0）: ").strip() or "0"

    checkpoint_path = f"checkpoints/{model_name}/fold_{fold}/best.pth"

    cmd = f"python baselines/{model_name}/infer.py --image {image_path} --checkpoint {checkpoint_path}"

    print(f"\n执行命令: {cmd}\n")
    os.system(cmd)


def infer_all_models():
    """所有模型推理"""
    image_path = input("请输入图像路径: ").strip()
    fold = input("使用哪个fold的模型（默认0）: ").strip() or "0"

    cmd = f"python infer_all.py --image {image_path} --fold {fold}"

    print(f"\n执行命令: {cmd}\n")
    os.system(cmd)


def show_project_structure():
    """显示项目结构"""
    print("\n项目结构：")
    print("""
radar_seg_project/
├── baselines/              # 所有baseline模型
│   ├── unet/              # UNet模型
│   ├── deeplabv3plus/     # DeepLabv3+模型
│   ├── fcn/               # FCN模型
│   ├── pspnet/            # PSPNet模型
│   ├── unetpp/            # UNet++模型
│   └── segnet/            # SegNet模型
├── datasets/              # 数据集加载
├── metrics/               # 评价指标
├── utils/                 # 工具函数
├── configs/               # 配置文件
├── data/                  # 数据目录
│   ├── images/           # 输入图像（1500张）
│   └── masks/            # 标签图像（1500张）
├── checkpoints/           # 模型权重保存
├── logs/                  # 训练日志
├── outputs/               # 推理结果输出
├── train_all.py          # 一键训练所有模型
├── infer_all.py          # 一键推理所有模型
└── README.md             # 使用说明
    """)


def show_usage():
    """显示使用说明"""
    print("\n使用说明：")
    print("""
1. 数据准备：
   - 将1500张输入图像放入 data/images/
   - 将1500张标签图像放入 data/masks/
   - 确保文件名一一对应

2. 训练模型：
   - 单个模型：选择菜单选项1或2
   - 所有模型：选择菜单选项3

3. 推理：
   - 训练完成后，使用菜单选项4或5进行推理
   - 结果保存在 outputs/<model_name>/ 目录

4. 查看结果：
   - 训练日志：logs/<model_name>/
   - 模型权重：checkpoints/<model_name>/fold_<x>/
   - 推理结果：outputs/<model_name>/

5. 评价指标：
   - IoU, mIoU, Dice, Pixel Accuracy, Precision, Recall
   - 结果保存在 logs/ 目录的JSON文件中

详细说明请查看 README.md
    """)


def main():
    """主函数"""
    while True:
        print_menu()
        choice = input("\n请输入选项编号: ").strip()

        if choice == '1':
            train_single_model_single_fold()
        elif choice == '2':
            train_single_model_all_folds()
        elif choice == '3':
            train_all_models()
        elif choice == '4':
            infer_single_model()
        elif choice == '5':
            infer_all_models()
        elif choice == '6':
            show_project_structure()
        elif choice == '7':
            show_usage()
        elif choice == '0':
            print("\n感谢使用！再见！\n")
            break
        else:
            print("\n无效选项，请重新输入")

        input("\n按回车键继续...")


if __name__ == '__main__':
    main()
